//
// Created by beich on 2020/11/9.
//

#include "linker_block_allocator.h"
#include <cstring>
#include <sys/mman.h>
#include <sys/prctl.h>
#include <cstdlib>

#if __ANDROID_API__ >= __ANDROID_API_Q__
static constexpr size_t kAllocateSize = PAGE_SIZE * 100;
#else
static constexpr size_t kAllocateSize = PAGE_SIZE;
#endif
static_assert(kAllocateSize % PAGE_SIZE == 0, "Invalid kAllocateSize.");

// the multiplier should be power of 2
static constexpr size_t round_up(size_t size, size_t multiplier) {
    return (size + (multiplier - 1)) & ~(multiplier - 1);
}

struct LinkerBlockAllocatorPage {
    LinkerBlockAllocatorPage *next;
    uint8_t bytes[kAllocateSize - 16] __attribute__((aligned(16)));
};

struct FreeBlockInfo {
    void *next_block;
    size_t num_free_blocks;
};

LinkerBlockAllocator::LinkerBlockAllocator(size_t block_size)
        : block_size_(round_up(block_size < sizeof(FreeBlockInfo) ? sizeof(FreeBlockInfo) : block_size, 16)),
          page_list_(nullptr),
          free_block_list_(nullptr) {
#if __ANDROID_API__ >= __ANDROID_API_Q__
    allocated_ = 0;
#endif
}

void *LinkerBlockAllocator::alloc() {
    protect_all(PROT_READ | PROT_WRITE);
    if (free_block_list_ == nullptr) {
        create_new_page();
    }

    FreeBlockInfo *block_info = reinterpret_cast<FreeBlockInfo *>(free_block_list_);
    if (block_info->num_free_blocks > 1) {
        FreeBlockInfo *next_block_info = reinterpret_cast<FreeBlockInfo *>(
                reinterpret_cast<char *>(free_block_list_) + block_size_);
        next_block_info->next_block = block_info->next_block;
        next_block_info->num_free_blocks = block_info->num_free_blocks - 1;
        free_block_list_ = next_block_info;
    } else {
        free_block_list_ = block_info->next_block;
    }

    memset(block_info, 0, block_size_);
#if __ANDROID_API__ >= __ANDROID_API_Q__
    ++allocated_;
#endif
    return block_info;
}

void LinkerBlockAllocator::free(void *block) {
    if (block == nullptr) {
        return;
    }
    protect_all(PROT_READ | PROT_WRITE);
    LinkerBlockAllocatorPage *page = find_page(block);
    if (page == nullptr) {
        abort();
    }

    ssize_t offset = reinterpret_cast<uint8_t *>(block) - page->bytes;

    if (offset % block_size_ != 0) {
        LOGW("The current page has a memory block that cannot be released: %p, offset: %zx, May cause memory leaks.", page, offset);
//		abort();
        return;
    }

    memset(block, 0, block_size_);

    FreeBlockInfo *block_info = reinterpret_cast<FreeBlockInfo *>(block);

    block_info->next_block = free_block_list_;
    block_info->num_free_blocks = 1;

    free_block_list_ = block_info;
#if __ANDROID_API__ >= __ANDROID_API_Q__
    --allocated_;
#endif
}

void LinkerBlockAllocator::protect_all(int prot) {
    for (LinkerBlockAllocatorPage *page = page_list_; page != nullptr; page = page->next) {
        if (mprotect(page, kAllocateSize, prot) == -1) {
            abort();
        }
    }
}

void LinkerBlockAllocator::create_new_page() {
    static_assert(sizeof(LinkerBlockAllocatorPage) == kAllocateSize,
                  "Invalid sizeof(LinkerBlockAllocatorPage)");

    LinkerBlockAllocatorPage *page = reinterpret_cast<LinkerBlockAllocatorPage *>(
            mmap(nullptr, kAllocateSize, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0));

    if (page == MAP_FAILED) {
        abort(); // oom
    }

    prctl(PR_SET_VMA, PR_SET_VMA_ANON_NAME, page, kAllocateSize, "linker_alloc");

    FreeBlockInfo *first_block = reinterpret_cast<FreeBlockInfo *>(page->bytes);
    first_block->next_block = free_block_list_;
    first_block->num_free_blocks = sizeof(page->bytes) / block_size_;

    free_block_list_ = first_block;

    page->next = page_list_;
    page_list_ = page;
}

LinkerBlockAllocatorPage *LinkerBlockAllocator::find_page(void *block) {
    if (block == nullptr) {
        abort();
    }

    LinkerBlockAllocatorPage *page = page_list_;
    while (page != nullptr) {
        const uint8_t *page_ptr = reinterpret_cast<const uint8_t *>(page);
        if (block >= (page_ptr + sizeof(page->next)) && block < (page_ptr + kAllocateSize)) {
            return page;
        }

        page = page->next;
    }

    abort();
}

void LinkerBlockAllocator::purge() {
#if __ANDROID_API__ >= __ANDROID_API_Q__
    if (allocated_) {
        return;
    }
#endif
    LinkerBlockAllocatorPage *page = page_list_;
    while (page) {
        LinkerBlockAllocatorPage *next = page->next;
        munmap(page, kAllocateSize);
        page = next;
    }
    page_list_ = nullptr;
    free_block_list_ = nullptr;
}
